"""
MongoDB 的数据持久层对象
@Time: 2021/2/19 14:11
@Author:
@File: mongo_po.py
"""
from pydantic import BaseModel, Field, validator
from datetime import datetime
from enum import Enum, unique
from typing import List, Optional
from bson.objectid import ObjectId



@unique
class ModelTypeEnum(Enum):
    NAMED_ENTITY = 'named entity'  # 命名实体训练
    REL_EXTRACTION = 'relation extraction'  # 关系抽取训练


class SampleDatasets(BaseModel):
    """对应于本次训练的数据集"""
    train_data: List[str] = Field(title='训练集，用来训练的样本的 id',
                                  description='可以根据该 id 去 ES 中该样本的全部信息',
                                  default_factory=list)
    validate_data: List[str] = Field(title='验证集，用来验证的样本的 id',
                                     description='可以根据该 id 去 ES 中该样本的全部信息',
                                     default_factory=list)
    test_data: List[str] = Field(title='测试集，用来测试的样本的 id',
                                 description='可以根据该 id 去 ES 中该样本的全部信息',
                                 default_factory=list)


class PerformanceIndex(BaseModel):
    """某个指标在不同数据集上的性能指标值"""
    train: Optional[float] = Field(None, title='在训练集上的指标值')
    valid: Optional[float] = Field(None, title='在验证集上的指标值')
    test: Optional[float] = Field(None, title='在测试集上的指标值')


class EpochStatus(BaseModel):
    """一轮迭代时的状态"""
    epoch_num: int = Field(title='迭代的轮次')
    acc: PerformanceIndex = Field(title='Accuracy',
                                  default_factory=PerformanceIndex)
    precision: PerformanceIndex = Field(title='Precision',
                                        default_factory=PerformanceIndex)
    recall: PerformanceIndex = Field(title='召回率',
                                     default_factory=PerformanceIndex)
    loss: PerformanceIndex = Field(title='损失值',
                                   default_factory=PerformanceIndex)
    

class MachineLearningModel(BaseModel):
    """
    样本训练的机器学习模型
    """
    _id: ObjectId = Field(title='文档的 ID',
                          default=None)
    train_name: str = Field(title='训练名称')
    version: str = Field(title='版本号',
                         default='0.0.1')
    created_time: datetime = Field(title='创建时间',
                                   default_factory=datetime.now)
    model_brief: str = Field(title='模型简介',
                             default='')
    model_type: str = Field(title='模型类型')
    datasets: SampleDatasets = Field(title='用来训练的样本的 id',
                                     description='可以根据该 id 去 ES 中该样本的全部信息',
                                     default_factory=SampleDatasets)
    train_status: int = Field(title='训练状态',
                              description='0 未训练, 1 训练中，2 已训练',
                              default=0)
    latest_train_time: datetime = Field(title='最近一次的训练时间',
                                        default=None)
    epochs_status: List[EpochStatus] = Field(title='所有轮迭代的指标',
                                             default_factory=list)
    hyper_params: dict = Field(title='模型的超参数',
                               default_factory=dict)
    train_params: dict = Field(title='模型的训练参数',
                               default_factory=dict)
    model_serialization: bytes = Field(title='模型序列化后的字节',
                                       default=b'')
    other_data: dict = Field(title='其他字段值，用于扩展',
                             default_factory=dict)
    parent_model: Optional[str] = Field(title='父模型的id',
                                        description='即该模型继承自哪个模型',
                                        default=None)
    comment: str = Field(title='备注',
                         default='')

    @validator('model_type')
    def model_type_match(cls, v):
        """验证`model_type`字段的值"""
        if v not in {'命名实体', '关系抽取'}:
            raise ValueError('模型类型必须为 [`命名实体`, `关系抽取`] 其中之一')
        return v

    def incr_version(self, inr_num: int = 1, loc: int = -1) -> str:
        """
        增加版本号
        例如原版本号为 0.0.0，调用该方法后则变为 0.0.1，再次调用并传入参数 (5, 0) 后则变为 5.0.1
        :param inr_num: 所要增加的数值，默认为1
        :param loc: 所要增加的位置，默认是第三位
        :return: 增加后的版本号
        """
        version_nums: List[str] = self.version.split('.')
        version_nums[loc] = str(int(version_nums[loc]) + inr_num)
        self.version = '.'.join(version_nums)
        return self.version
